﻿#region MIT License
/*

Copyright (C) 2011 by Timothy A. Barela

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.

*/
#endregion

using System;
using System.Collections.Generic;
using System.Diagnostics;
using Mono.Cecil;
using Mono.Cecil.Cil;
using Mono.Cecil.Rocks;

namespace Sharpener
{
	/// <summary>
	/// Caches and injects code used for method inlining.
	/// </summary>
	internal class InlineMethod
	{
		static Type forceInlineAttributeType;

		public MethodDefinition SourceMethod { get; private set; }
		List<Instruction> instructions;

		static InlineMethod()
		{
			forceInlineAttributeType = typeof(ForceInlineAttribute);
		}

		public InlineMethod( MethodDefinition method )
		{
			SourceMethod = method;
						
			method.Body.SimplifyMacros();
			instructions = InstructionHelper.CloneList(method.Body.Instructions);
			method.Body.OptimizeMacros();
		}
		
		public void Inject( MethodDefinition dstMethod, int index )
		{
			//create "working copy" of the target method
			var workInstructions = new List<Instruction>(dstMethod.Body.Instructions);
			
			//clone this patch's instructions-- these will end up in the final "product".
			var patchedInstructions = InstructionHelper.CloneList(instructions);

			var callInstruction = workInstructions[index];
			var returnBranchTarget = workInstructions[index + 1];
								
			//-----------------------------
						
			//parameters
			var parameterMapping = new ParameterMapping(SourceMethod,dstMethod);
			parameterMapping.Remap(patchedInstructions);
			
			//local variables
			var localMapping = new LocalMapping(SourceMethod,dstMethod);
			localMapping.Remap(patchedInstructions);

			storeParameterLocals(patchedInstructions,parameterMapping);

			//can't leave returns from callee in place...
			convertReturnsToBranches(patchedInstructions,returnBranchTarget);

			//insert the patch into work.. ( patch-local branches are already corrected during previous clone ) 
			workInstructions.InsertRange(index+1,patchedInstructions);

			//-----------------------------

			//remove call instruction
			workInstructions.Remove(callInstruction);
			//...note that when we add inlining to methods with instance + parameters,
			//we'll have to reexamine where to remap to!!
			InstructionHelper.RemapBranchTargets(workInstructions,callInstruction,patchedInstructions[0]);

			//write instructions back to caller 
			overwriteTargetMethod(dstMethod, workInstructions);
		}

		private void convertReturnsToBranches( List<Instruction> patch, Instruction branchTarget )
		{
			for(int i = 0; i < patch.Count; i++)
			{
				if(patch[i].OpCode == OpCodes.Ret)
				{
					patch[i].OpCode = OpCodes.Br;
					patch[i].Operand = branchTarget;
				}
			}
		}

		private void overwriteTargetMethod( MethodDefinition method, List<Instruction> work )
		{
			method.Body.Instructions.Clear();

			foreach(var instruction in work)
			{
				method.Body.Instructions.Add(instruction);
			}
		}

		//load our param mapped locals now, if any...
		//cheap, but at least we are inlining parameterized methods now!
		private void storeParameterLocals(IList<Instruction> instructions,ParameterMapping paramMapping)
		{
			if( paramMapping.Count > 0 )
			{
				foreach(var srcParam in paramMapping.SourceMethod.Parameters)
				{
					var localVar = paramMapping[srcParam] as VariableDefinition;
					var inst = Instruction.Create(OpCodes.Stloc,localVar);
					instructions.Insert(0,inst);
				}
			}
		}
		
		internal static Dictionary<string,InlineMethod> FindAll(AssemblyDefinition assembly)
		{
			var inlinePatches = new Dictionary<string,InlineMethod>();

			foreach(var module in assembly.Modules)
			{
				foreach(var type in module.Types)
				{
					foreach(var method in type.Methods)
					{
						if(isMethodInlineable(method))
						{
							inlinePatches.Add(method.FullName,new InlineMethod(method));
						}
					}
				}
			}
			
			return inlinePatches;
		}
		
		//Check to see if method meets our restrictions on what can be inlined. 
		private static bool isMethodInlineable(MethodDefinition method)
		{
			if(!method.HasBody)
				return false;

			if(method.IsNative)
				return false;

			if(method.IsVirtual)
				return false;

			if(method.IsConstructor)
				return false;

			if(method.HasGenericParameters)
				return false;

			//the following are restrictions imposed during development...
			//we plan to remove these over time.

			if(!method.IsStatic)
				return false;
			
			if(method.HasCustomAttributes)
			{
				foreach(var ca in method.CustomAttributes)
				{
					if(ca.AttributeType.FullName == forceInlineAttributeType.FullName)
					{
						return true;
					}
				}
			}

			return false;
		}

		internal void CleanUp()
		{
			var method = SourceMethod;

			for(int i = 0;i < method.CustomAttributes.Count;i++)
			{
				var ca = method.CustomAttributes[i];

				if(ca.AttributeType.FullName == forceInlineAttributeType.FullName)
				{
					//yes this is hardcoded magic...
					//gets the RemoveFromAssembly property.
					//**KLUDGE** for some reason, cecil wont notice the property
					//unless it was explicitly set as a named parameter(?!?)
					//..so we have to check if there is any properties at all
					var result = ca.Properties.Count > 0 ? (bool)ca.Properties[0].Argument.Value : true;

					if(result)
					{
						Debug.Print("Removing method '{0}' from output assembly.",method.FullName);
						method.DeclaringType.Methods.Remove(method);
					}
					else
					{
						//..just remove the attribute then
						Debug.Print("Removing ForceInlineAttribute from method {0}",method.FullName);
						method.CustomAttributes.RemoveAt(i);
					}

					return;
				}
			}

			//if we ever get here, something has gone very wrong.
			throw new Exception("Expected Sharpener.ForceInlineAttribute.RemoveFromAssembly, " +
						"but no matching property was found.");
		}
	}
}